# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jaewon

import logging

from coin.base.timestamp import get_timestamp
from coin.exchange.base.order_gateway import OrderSide, OrderType
from coin.exchange.bitmex.kr_rest.futures_product import BitmexFuturesProduct


class BitmexSimpleMMExecutor(object):
  def __init__(self, execution_params, order_gateway, logger=None):
    self._logger = logger or logging.getLogger('BitmexSimpleMMExecutor')
    self._og = order_gateway

    self._product = execution_params['product']
    self._native_symbol = execution_params['product'].native_symbol

    self._lot_size = execution_params['lot_size']
    self._stickiness = execution_params['stickiness']
    self._min_pos = execution_params['min_pos']
    self._max_pos = execution_params['max_pos']
    self._post_only = execution_params.get('post_only', True)
    self._order_update_period = (execution_params.get('order_update_period', 2.) * (10**9))

    self._pos = 0
    self._sell_on_the_fly_orders = []
    self._sell_working_orders = []
    self._buy_on_the_fly_orders = []
    self._buy_working_orders = []

    self._last_fire_time = 0
    self._create_list = []
    self._amend_list = []

  def _update_orders(self):
    self._pos = self._og.get_position(self._product)

    self._sell_on_the_fly_orders = []
    self._sell_working_orders = []
    self._buy_on_the_fly_orders = []
    self._buy_working_orders = []

    for order in self._og.get_working_order():
      if (order.product != self._product or not order.order_id.startswith('pass-')):
        continue

      if order.order_side == OrderSide.SELL:
        if order.accepted:
          self._sell_working_orders.append(order)
        else:
          self._sell_on_the_fly_orders.append(order)

      elif order.order_side == OrderSide.BUY:
        if order.accepted:
          self._buy_working_orders.append(order)
        else:
          self._buy_on_the_fly_orders.append(order)

      else:
        self._logger.error('Unknown side: %s', str(order))
        continue

  def _cancel_sell_working_orders(self, exclude_order_id=None):
    for order in self._sell_working_orders:
      if order.order_id == exclude_order_id:
        continue
      self._amend_list.append(self._og.make_bitmex_amend(order.order_id, leaves_qty=0))

  def _cancel_buy_working_orders(self, exclude_order_id=None):
    for order in self._buy_working_orders:
      if order.order_id == exclude_order_id:
        continue
      self._amend_list.append(self._og.make_bitmex_amend(order.order_id, leaves_qty=0))

  def _manage_sell_order(self, sell_price):
    # If there is an on-the-fly order, cancel all working orders and keep the
    # on-the-fly order.
    if (sell_price is None or self._sell_on_the_fly_orders or self._pos <= self._min_pos):
      self._cancel_sell_working_orders()
      return

    assert self._pos - self._min_pos > 0, self._pos
    sell_qty = min(self._pos - self._min_pos, self._lot_size)
    assert sell_qty > 0, sell_qty

    # Create new order if there is no working order.
    if not self._sell_working_orders:
      order_id = 'pass-%d' % get_timestamp()
      order = self._og.make_bitmex_order(product=self._product,
                                         price=sell_price,
                                         qty=sell_qty,
                                         order_side=OrderSide.SELL,
                                         order_type=OrderType.LIMIT,
                                         order_id=order_id,
                                         post_only=self._post_only,
                                         tag='pass')
      self._create_list.append(order)
      self._logger.info('SELL %s %f @ %.1f, tag=pass', self._native_symbol, sell_qty, sell_price)
      return

    # If there is an working order that its price is same to sell_price,
    # keep it and cancel all other orders.
    for order in self._sell_working_orders:
      if abs(order.price - sell_price) <= self._stickiness:
        self._cancel_sell_working_orders(exclude_order_id=order.order_id)
        return

    # Choose one with the smallest price and modify it and
    # cancel all other orders.
    for idx, order in enumerate(self._sell_working_orders):
      if idx == 0:
        chosen_order = order
        continue
      if order.price < chosen_order.price:
        chosen_order = order

    self._cancel_sell_working_orders(exclude_order_id=chosen_order.order_id)
    self._amend_list.append(
        self._og.make_bitmex_amend(chosen_order.order_id, price=sell_price, leaves_qty=sell_qty))
    self._logger.info('SELL %s %f @ %.1f, tag=pass', self._native_symbol, sell_qty, sell_price)

  def _manage_buy_order(self, buy_price):
    # If there is an on-the-fly order, cancel all working orders and keep the
    # on-the-fly order.
    if (buy_price is None or self._buy_on_the_fly_orders or self._pos >= self._max_pos):
      self._cancel_buy_working_orders()
      return

    assert self._max_pos - self._pos > 0, self._pos
    buy_qty = min(self._max_pos - self._pos, self._lot_size)
    assert buy_qty > 0, buy_qty

    # Create new order if there is no working order.
    if not self._buy_working_orders:
      order_id = 'pass-%d' % get_timestamp()
      order = self._og.make_bitmex_order(product=self._product,
                                         price=buy_price,
                                         qty=buy_qty,
                                         order_side=OrderSide.BUY,
                                         order_type=OrderType.LIMIT,
                                         order_id=order_id,
                                         post_only=self._post_only,
                                         tag='pass')
      self._create_list.append(order)
      self._logger.info('BUY  %s %f @ %.1f, tag=pass', self._native_symbol, buy_qty, buy_price)
      return

    # If there is an working order that its price is same to buy_price,
    # keep it and cancel all other orders.
    for order in self._buy_working_orders:
      if abs(order.price - buy_price) <= self._stickiness:
        self._cancel_buy_working_orders(exclude_order_id=order.order_id)
        return

    # Choose one with the biggest price and modify it and
    # cancel all other orders.
    for idx, order in enumerate(self._buy_working_orders):
      if idx == 0:
        chosen_order = order
        continue
      if order.price > chosen_order.price:
        chosen_order = order

    self._cancel_buy_working_orders(exclude_order_id=chosen_order.order_id)
    self._amend_list.append(
        self._og.make_bitmex_amend(chosen_order.order_id, price=buy_price, leaves_qty=buy_qty))
    self._logger.info('BUY  %s %f @ %.1f, tag=pass', self._native_symbol, buy_qty, buy_price)

  def execute(self, create_list, amend_list):
    if get_timestamp() - self._last_fire_time < self._order_update_period:
      return

    if amend_list:
      self._og.amend_bulk(amend_list)
      self._logger.debug('Amend %d orders', len(amend_list))

    if create_list:
      self._og.submit_bulk(create_list)
      self._logger.debug('Create %d orders', len(create_list))

    if create_list or amend_list:
      self._last_fire_time = get_timestamp()

  def manage_mm_orders(self, sell_price, buy_price, create_list=None, amend_list=None, fire=True):
    self._create_list = [] if create_list is None else create_list
    self._amend_list = [] if amend_list is None else amend_list
    self._update_orders()

    self._manage_sell_order(sell_price)
    self._manage_buy_order(buy_price)

    if fire:
      self.execute(self._create_list, self._amend_list)
